import httpx
import requests
import json
import logging
from PyCmpltrtok.common import md5, sep
from python_basic.exercise.lib.log.learn_from_fastchat.logger_util import build_logger

logger = build_logger('fastapi_interface', 'fastapi_interface.log', logger_level=logging.INFO)


def text_gen(xinput, username='', dport=7750, timeout=99999999):
    xjson = dict()
    xjson['xinput'] = xinput
    xjson['username'] = username
    # xjson = dict(query=xjson)

    base_url = f'http://127.0.0.1:{dport}'
    print(base_url)
    url = '/stream_chat'
    print(url)
    entire_url = f'{base_url}{url}'
    print(entire_url)
    
    res = requests.post(entire_url, json=xjson, stream=True)
    for chunk in res.iter_lines(decode_unicode=False, delimiter=b'\0'):
        if not chunk:
            continue
        chunk = chunk.decode('utf8')
        data = json.loads(chunk)
        xouput = data['output']
        yield xouput
    
    # _client = httpx.Client(base_url=base_url, timeout=timeout)
    # response = _client.stream("POST", url, json=xjson, data=None)
    # as_json = True
    
    # try:
    #     with response as r:
    #         for chunk in r.iter_text(None):
    #             if not chunk: # fastchat api yield empty bytes on start and end
    #                 continue
    #             if as_json:
    #                 try:
    #                     if chunk.startswith("data: "):
    #                         data = json.loads(chunk[6:-2])
    #                     elif chunk.startswith(":"): # skip sse comment line
    #                         continue
    #                     else:
    #                         data = json.loads(chunk)
    #                     yield data['output']
    #                 except Exception as e:
    #                     msg = f"接口返回json错误： ‘{chunk}’。错误信息是：{e}。"
    #                     logger.error(f'{e.__class__.__name__}: {msg}', exc_info=e)
    #             else:
    #                 # print(chunk, end="", flush=True)
    #                 yield chunk
    # except httpx.ConnectError as e:
    #     msg = f"无法连接API服务器，请确认 ‘api.py’ 已正常启动。({e})"
    #     logger.error(msg)
    #     yield {"code": 500, "msg": msg}
    # except httpx.ReadTimeout as e:
    #     msg = f"API通信超时，请确认已启动FastChat与API服务（详见Wiki '5. 启动 API 服务或 Web UI'）。（{e}）"
    #     logger.error(msg)
    #     yield {"code": 500, "msg": msg}
    # except Exception as e:
    #     msg = f"API通信遇到错误：{e}"
    #     logger.error(f'{e.__class__.__name__}: {msg}', exc_info=e)
    #     yield {"code": 500, "msg": msg}


if '__main__' == __name__:
    user = 'u0001'
    gen = text_gen('怎么滑雪？', username=user)
    for x in gen:
        print(x)
        